import struct
import numpy as np
from PIL import Image
import os

def load_mnist_images(filename):
    """读取 MNIST 图像"""
    with open(filename, 'rb') as f:
        magic, num_images, rows, cols = struct.unpack(">IIII", f.read(16))
        images = np.frombuffer(f.read(), dtype=np.uint8).reshape(num_images, rows, cols)
    return images

def load_mnist_labels(filename):
    """读取 MNIST 标签"""
    with open(filename, 'rb') as f:
        magic, num_labels = struct.unpack(">II", f.read(8))
        labels = np.frombuffer(f.read(), dtype=np.uint8)
    return labels

def save_images_by_label(images, labels, output_dir):
    """按标签分类保存图像到文件夹"""
    os.makedirs(output_dir, exist_ok=True)
    for idx, (img, lbl) in enumerate(zip(images, labels)):
        label_dir = os.path.join(output_dir, str(lbl))
        os.makedirs(label_dir, exist_ok=True)
        img_path = os.path.join(label_dir, f"{idx}.png")
        Image.fromarray(img).save(img_path)
    print(f"已保存 {len(images)} 张图片到 {output_dir}")

if __name__ == "__main__":
    data_folder = "../data/MNIST/raw"  # MNIST 文件所在目录

    datasets = [
        ("train-images-idx3-ubyte", "train-labels-idx1-ubyte", "train"),
        ("t10k-images-idx3-ubyte", "t10k-labels-idx1-ubyte", "test")
    ]

    for img_file, lbl_file, name in datasets:
        print(f"处理 {name} 集合...")
        images = load_mnist_images(os.path.join(data_folder, img_file))
        labels = load_mnist_labels(os.path.join(data_folder, lbl_file))
        save_images_by_label(images, labels, os.path.join(data_folder, f"{name}_by_label"))
